import time
from mylib import UartCommend, Camera
import tf

# 初始化串口通信
uartcommend = UartCommend(uart_port=1, baudrate=115200)

# 初始化摄像头
camera = Camera()

# 模型和标签文件路径（SD卡根目录）
model_path = "/sd/model.tflite"
label_path = "/sd/label.txt"

# 加载模型和标签
print(f"Loading model: {model_path}")
tf_model = tf.load(model_path)

print(f"Loading labels: {label_path}")
with open(label_path, "r") as f:
    labels = [line.strip() for line in f]


# 定义小类和大类
small_classes = {
    10: 'keyboard',
    11: 'mobile_phone',
    12: 'mouse',
    13: 'headphones',
    14: 'monitor',
    15: 'speaker',
    1: 'wrench',
    2: 'soldering_iron',
    3: 'electrodrill',
    4: 'tape_measure',
    5: 'screwdriver',
    6: 'pliers',
    7: 'oscillograph',
    8: 'multimeter',
    9: 'printer'
}

# 大类映射
electronics = {'keyboard', 'mouse', 'headphones', 'monitor', 'speaker', 'printer', 'mobile_phone'}
tools = {'wrench', 'soldering_iron', 'electrodrill', 'tape_measure', 'screwdriver', 'pliers', 'oscillograph', 'multimeter'}

def run_model(image):
    """
    在图像上运行模型并返回分类结果。如果置信度大于0.7且在10次检测中结果一致，则返回大类分类（left or right）。
    """
    consistent_results = []  # 用于保存连续10次的分类结果
    max_attempts = 10  # 最大检测次数

    for _ in range(max_attempts):
        for obj in tf.classify(tf_model, image, min_scale=1.0, scale_mul=0.5, x_overlap=-1, y_overlap=-1):
            # 获取分类结果并排序
            sorted_list = sorted(zip(labels, obj.output()), key=lambda x: x[1], reverse=True)
            top_result = sorted_list[0]  # 获取最可信的结果
            
            label = top_result[0]
            confidence = top_result[1]
            
            # 如果置信度大于0.7
            if confidence > 0.7:
                # 检查小类对应的大类
                if label in electronics:
                    result =  0x01 #'left'
                elif label in tools:
                    result =  0x02 #'right'
                else:
                    result = None
                
                # 将结果加入到检测结果列表
                consistent_results.append(result)

                # 如果检测结果有10次，检查是否一致
                if len(consistent_results) == max_attempts:
                    if all(res == consistent_results[0] for res in consistent_results):
                        return consistent_results[0]  # 如果一致，返回结果
                    else:
                        return None  # 如果不一致，返回 None

    # 如果没有返回一致的结果，返回 None
    return None


while True:
    # 获取图像和FPS
    img_roi = camera.get_image()

    # 处理接收到的串口数据
    while uartcommend.uart.any():
        data = uartcommend.get_uart_data()
        if data:
            if data[0] == 0x01:  # 如果数据是0x01，运行模型
                # 跑模型
                result = run_model(img_roi)
                if result:
                    uartcommend.send_uart_data(result)  # 发送确认消息 撞击
                else:
                    print("No detection")

